# ~*~ coding: utf-8 ~*~
import re

import psutil
import time
import os

from flask import current_app
from concurrent.futures import ProcessPoolExecutor
from celery.exceptions import SoftTimeLimitExceeded

from apps.extensions import celery, db
from apps.tasks.models import PeriodicTask, PeriodicTaskRunHistory
from apps.ops.celery.decorator import after_app_ready_start
from apps.ops.inventory import BeesCustomInventory
from apps.ops.ansible.runner import AdHocRunner, PlayBookRunner

from apps.config import ConfigManager

CONFIG = ConfigManager.load_user_config()


@celery.task
def run_periodic_task(task_id):
    task = PeriodicTask.query.get(task_id)

    # 该任务还有没有执行完毕的需要终止此次执行
    if task.only_once and task.has_running:
        current_app.logger.warning('The current task({0}) has enabled the "only once" and has terminated running'.
                                   format(task.name))
        return

    # 运行该任务
    try:
        task.run()
    except SoftTimeLimitExceeded:
        pass


@celery.task
@after_app_ready_start
def process_terminate_task():
    histories = db.session.query(PeriodicTaskRunHistory).\
        filter(PeriodicTaskRunHistory.is_finished == 2).\
        all()
    if len(histories) > 0:
        pool = ProcessPoolExecutor(len(histories))
        for history in histories:
            if not history.process_id:
                continue
            pool.submit(hold_process, history)
        pool.shutdown(wait=True)


def hold_process(history):
    history = db.session.merge(history)
    try:
        p = psutil.Process(history.process_id)
        p.wait()
    except psutil.NoSuchProcess:
        pass

    log_absolute_path = os.path.join(CONFIG.TASKS_LOG_DIR, history.log_path)
    try:
        with open(log_absolute_path, 'r+', encoding='utf-8') as f:
            text = f.read()
            response = re.findall("""\{'data': '(.*)', 'is_success': (.*)\}""", text)[0]
        result = {'data': response[0], 'is_success': True if eval(response[1]) else False}
    except IndexError:  # 异常终止, 比如进程被killed
        result = {'data': 'Terminated', 'is_success': 0}

    # 处理执行结果
    if result['data'] == 'Terminated':
        history.is_finished = 0
    else:
        history.is_finished = 1
    history.result = result['data']
    history.is_success = result['is_success']
    # 处理计算时间
    time_end = time.time()
    history.date_finished = time_end
    history.timedelta = time_end - history.date_start
    history.update()


@celery.task
def run_adhoc(hosts, groups, tasks, log):
    inventory = BeesCustomInventory(hosts, groups)
    runner = AdHocRunner(inventory)
    response = runner.run(tasks, "all")
    result = {
        'raw': response.results_raw,
        'summary': response.results_summary
    }

    log = db.session.merge(log)
    log.update(result=result)


@celery.task
def run_playbook(hosts, groups, playbook_path, log):
    inventory = BeesCustomInventory(hosts, groups)
    runner = PlayBookRunner(inventory, options={'playbook_path': playbook_path})
    result = runner.run()

    log = db.session.merge(log)
    log.update(result=result.get('stats'))


@celery.task
def run_role(hosts, groups, role_path, log):
    inventory = BeesCustomInventory(hosts, groups)
    runner = PlayBookRunner(inventory, options={'playbook_path': role_path})
    result = runner.run()

    log = db.session.merge(log)
    log.update(result=result.get('stats'))
